import base64
import hashlib
import os
import time
from typing import Union, List, Dict
from PIL import Image
from io import BytesIO
from langchain_openai.chat_models.base import BaseChatOpenAI
from loguru import logger


def generate_image_file(llm: BaseChatOpenAI,
                        image_desc: str,
                        image_size: Union[str, List[str]] = "1024x1024",
                        image_type: str = "png") -> Union[str, Dict[str, str], None]:
    try:
        logger.debug(
            f"Generating image file - desc: '{image_desc[:30]}...', size: {image_size}, type: {image_type}")
        valid_image_sizes = llm.model_config.get("valid_size", [])

        image_base64_result = generate_image_base64(llm, image_desc, image_size, valid_image_sizes)
        if not image_base64_result:
            logger.warning(f"Failed to generate base64 image for: '{image_desc[:30]}...'")
            return None

        def save_single_image(base64_str: str, size: str = "") -> str:
            timestamp = int(time.time() * 1000)
            filename = f"{timestamp}_{size}.{image_type}".strip("_")
            filepath = os.path.join(get_or_create_dir(filename=filename), filename)
            img_data = base64.b64decode(base64_str)
            with open(filepath, 'wb') as f:
                f.write(img_data)
            logger.debug(f"Image saved{' (size: ' + size + ')' if size else ''}: {filepath}")
            return filename

        if isinstance(image_base64_result, str):
            size = image_size[0] if isinstance(image_size, list) else image_size
            return {size: save_single_image(image_base64_result, size)}
        else:
            return {size: save_single_image(base64_str, size)
                    for size, base64_str in image_base64_result.items()}

    except Exception as e:
        logger.error(f"Image generation failed: {str(e)}")
        return None


def generate_image_base64(llm: BaseChatOpenAI,
                          image_desc: str,
                          image_size: Union[str, List[str]] = "1024x1024",
                          valid_image_sizes: List[str] = []) -> Union[str, Dict[str, str], None]:
    if not llm:
        logger.error("No LLM provided for image generation")
        return None

    try:
        image_sizes = [image_size] if isinstance(image_size, str) else image_size

        target_size = max(image_sizes, key=lambda x: int(x.split('x')[0]) * int(x.split('x')[1]))

        if valid_image_sizes:
            # Parse target dimensions
            target_w, target_h = map(int, target_size.split('x'))
            target_ratio = target_w / target_h

            # First check for exact match
            if target_size in valid_image_sizes:
                target_size = target_size
            else:
                # Look for same aspect ratio
                same_ratio_sizes = []
                for size in valid_image_sizes:
                    w, h = map(int, size.split('x'))
                    ratio = w / h
                    if abs(ratio - target_ratio) < 0.001:  # Account for float precision
                        same_ratio_sizes.append(size)

                if same_ratio_sizes:
                    # Pick the closest resolution
                    target_size = min(same_ratio_sizes,
                                      key=lambda x: abs(int(x.split('x')[0]) - target_w))
                else:
                    # Pick size with closest ratio
                    target_size = min(valid_image_sizes,
                                      key=lambda x: abs(int(x.split('x')[0]) / int(x.split('x')[1]) - target_ratio))

        result = llm.root_client.images.generate(
            model=llm.model_name,
            prompt=image_desc,
            size=target_size,
            response_format="b64_json"
        )

        base_output = result.data[0].b64_json if hasattr(result, 'data') else None
        if not base_output:
            logger.warning("Image generation returned no base64 data")
            return None

        # If only one size requested, return string
        if len(image_sizes) == 1:
            if target_size != image_sizes[0]:
                logger.debug(f"Resizing from {target_size} to {image_sizes[0]}")
                return resize_base64_image(base_output, image_sizes[0])
            return base_output

        # Generate all requested sizes
        outputs = {}
        for size in image_sizes:
            if size == target_size:
                outputs[size] = base_output
            else:
                logger.debug(f"Resizing from {target_size} to {size}")
                outputs[size] = resize_base64_image(base_output, size)
        return outputs

    except Exception as e:
        logger.error(f"Image generation error: {str(e)}")
        return None


def resize_base64_image(base64_str: str, target_size: str, keep_ratio: bool = False) -> str:
    """
    Resize a base64 encoded image to the target size

    Args:
        base64_str: Base64 encoded image string
        target_size: Target size in format "widthxheight"
        keep_ratio: Whether to maintain aspect ratio

    Returns:
        Base64 encoded resized image string
    """
    try:
        # Decode base64 to image
        img_data = base64.b64decode(base64_str)
        img = Image.open(BytesIO(img_data))

        # Resize with high-quality compression
        target_w, target_h = map(int, target_size.split('x'))

        if keep_ratio:
            # Maintain aspect ratio using thumbnail method
            img.thumbnail((target_w, target_h), Image.Resampling.LANCZOS)
        else:
            # Force resize to exact dimensions
            img = img.resize((target_w, target_h), Image.Resampling.LANCZOS)

        # Convert back to base64
        buffered = BytesIO()
        img.save(buffered, format=img.format or "PNG")
        return base64.b64encode(buffered.getvalue()).decode()
    except Exception as e:
        logger.error(f"Image resize error: {str(e)}")
        return base64_str


def get_base64_file_size(base64_str) -> int:
    """
    Get the size in bytes of a base64 encoded file

    Args:
        base64_str: Base64 encoded string

    Returns:
        Size in bytes
    """
    if ',' in base64_str:
        base64_str = base64_str.split(',')[1]

    padding = len(base64_str) % 4
    if padding:
        base64_str += '=' * (4 - padding)
    return len(base64.b64decode(base64_str))


def get_or_create_dir(filename: str) -> str:
    """
    Get or create a directory path based on filename hash

    Args:
        filename: Filename to hash

    Returns:
        Directory path
    """
    h = md5_encode(filename)
    root_path = os.getenv("TMP_PATH", "/tmp")
    path = f"{root_path}/{h[1:3]}/{h[4:6]}"
    if not os.path.exists(path):
        os.makedirs(path)
    return path


def md5_encode(string) -> str:
    """
    Create MD5 hash of a string

    Args:
        string: String to hash

    Returns:
        MD5 hash as hexadecimal string
    """
    md5 = hashlib.md5()
    md5.update(string.encode('utf-8'))
    hex_digest = md5.hexdigest()
    return hex_digest
